from llama_index.core import VectorStoreIndex, Document, get_response_synthesizer
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from pprint import pprint

# 创建示例文档
sample_texts = [
    "LlamaIndex提供了多种检索策略，包括向量检索、关键词检索等。",
    "相似度阈值可以过滤掉相关性较低的检索结果，提高回答质量。",
    "自定义查询引擎允许开发者精确控制检索和响应过程。",
    "后处理器可以对检索结果进行进一步筛选和优化。",
]
documents = [Document(text=text) for text in sample_texts]
index = VectorStoreIndex.from_documents(documents)

# 配置检索器
retriever = VectorIndexRetriever(index=index, similarity_top_k=10)
# 配置响应合成器  响应合成器用来将检索到的节点信息整合成最终的答案
response_synthesize = get_response_synthesizer()
# 配置相似度后处理器 只保留相似分类大于。7的节点，过滤掉相关性比较低的结果节点
similarityPostprocessor = SimilarityPostprocessor(similarity_cutoff=0.8)
# 组装自定义的查询引擎
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesize,
    node_postprocessors=[similarityPostprocessor],  # 应用后处理器
)
response = query_engine.query("如何提高检索结果的质量?")
pprint(response)

print(f"检索到的节点数量:{len(response.source_nodes)}")
for i, node in enumerate(response.source_nodes):
    print(f"节点 {i+1} : {node.text[:100]}")
